#!/usr/bin/env python3
"""
Base64文本转PNG图片脚本
"""

import sys
import base64
from PIL import Image
import io

# 默认文件路径
base64_path = "seg_base64.txt" if len(sys.argv) < 2 else sys.argv[1]
output_path = "seg_result_mask.png" if len(sys.argv) < 3 else sys.argv[2]

# 读取base64文本
with open(base64_path, 'r') as f:
    base64_str = f.read().strip()

# 移除数据URI前缀（如果有）
if ',' in base64_str:
    base64_str = base64_str.split(',')[1]

# 解码base64
try:
    image_data = base64.b64decode(base64_str)
except Exception as e:
    print(f"❌ Base64解码失败: {e}")
    sys.exit(1)

# 转换为图片
try:
    image = Image.open(io.BytesIO(image_data))
    image.save(output_path, 'PNG')
    print(f"✅ 图片已保存: {output_path}")
    print(f"   尺寸: {image.size}")
except Exception as e:
    print(f"❌ 图片处理失败: {e}")
    sys.exit(1)

